# encoding utf-8

import requests
import pandas as pd
from datetime import datetime, timedelta
from typing import Optional
from config.settings import settings
from loguru import logger


class StockMarginService:
    """Stock Margin Trading Service"""
    
    def __init__(self):
        self.api_base_url = f"{settings.data_query_host}/data_query/stock_margin"
    
    def get_margin_balance(self, exchange: str = "all",
                          start_date: Optional[str] = None, end_date: Optional[str] = None) -> pd.DataFrame:
        """
        Query margin trading balance data.
        
        Args:
            exchange (str): Exchange type, supports 'SSE' (Shanghai), 'SZSE' (Shenzhen), 'all'.
            start_date (str, optional): Start date in "YYYY-MM-DD" format.
            end_date (str, optional): End date in "YYYY-MM-DD" format.
            
        Returns:
            pd.DataFrame: DataFrame containing margin trading data.
            
        Raises:
            ValueError: If the parameters are incorrect.
        """
        # Normalize exchange parameter
        exchange_mapping = {
            'sz': 'SZSE',
            'sh': 'SSE',
            'szse': 'SZSE', 
            'sse': 'SSE',
            'all': 'all'
        }
        
        normalized_exchange = exchange_mapping.get(exchange.lower(), exchange)
        if normalized_exchange not in ['SSE', 'SZSE', 'all']:
            raise ValueError("Exchange parameter must be 'SSE', 'SZSE', or 'all'")
        
        # Normalize dates to YYYY-MM-DD format
        start_date, end_date = self._normalize_dates(start_date, end_date)
        
        try:
            # Call API to get margin data
            data = self._fetch_margin_data_from_api(normalized_exchange, start_date, end_date)
            
            # Convert to DataFrame
            result = self._convert_api_response_to_dataframe(data)
            
            return result
            
        except Exception as e:
            logger.error(f"Failed to fetch margin data: {e}")
            return pd.DataFrame()
    
    def _fetch_margin_data_from_api(self, exchange: str, start_date: str, end_date: str) -> dict:
        """从API获取融资融券数据"""
        params = {
            'start_date': start_date,
            'end_date': end_date,
            'exchange': exchange
        }
        
        try:
            response = requests.get(
                self.api_base_url,
                params=params,
                timeout=settings.request_timeout_seconds,
            )
            response.raise_for_status()
            return response.json()
        except requests.exceptions.RequestException as e:
            logger.error(f"API request failed: {e}")
            raise
        except Exception as e:
            logger.error(f"Failed to parse API response: {e}")
            raise

    def _convert_api_response_to_dataframe(self, api_response: dict) -> pd.DataFrame:
        """将API响应转换为DataFrame"""
        if api_response.get('code') != 200:
            logger.error(f"API returned error: {api_response.get('message', 'Unknown error')}")
            return pd.DataFrame()
        
        data = api_response.get('data', [])
        if not data:
            logger.warning("API returned empty data")
            return pd.DataFrame()
        
        # Convert to DataFrame
        df = pd.DataFrame(data)
        
        # Ensure correct column names and data types
        expected_columns = {
            'date': 'date',
            'exchange_name': 'exchange_name', 
            'financing_balance': 'financing_balance',
            'financing_buy_amount': 'financing_buy_amount',
            'securities_lending_volume': 'securities_lending_volume',
            'securities_lending_balance': 'securities_lending_balance',
            'securities_lending_sell_volume': 'securities_lending_sell_volume',
            'total_margin_balance': 'total_margin_balance'
        }
        
        # Ensure all expected columns exist
        for col in expected_columns:
            if col not in df.columns:
                df[col] = None
        
        # Convert numeric columns
        numeric_columns = [
            'financing_balance', 'financing_buy_amount', 'securities_lending_volume',
            'securities_lending_balance', 'securities_lending_sell_volume', 'total_margin_balance'
        ]
        
        for col in numeric_columns:
            if col in df.columns:
                df[col] = pd.to_numeric(df[col], errors='coerce')
        
        # Reorder columns
        column_order = list(expected_columns.keys())
        df = df[column_order]
        
        return df

    
    def _normalize_dates(self, start_date, end_date):
        """Normalize start and end dates to YYYY-MM-DD format."""
        def format_date(date_str):
            if not date_str:
                return datetime.now().strftime("%Y-%m-%d")
            
            # If already in YYYY-MM-DD format
            if len(date_str) == 10 and date_str.count('-') == 2:
                return date_str
            
            # If in YYYYMMDD format
            if len(date_str) == 8 and date_str.isdigit():
                return f"{date_str[:4]}-{date_str[4:6]}-{date_str[6:8]}"
            
            return date_str
        
        if not start_date:
            current_date = datetime.now().strftime("%Y-%m-%d")
            start_date = current_date
            end_date = end_date or current_date
        elif not end_date:
            end_date = start_date
        
        return format_date(start_date), format_date(end_date)
    
    
    

